{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "from glob import glob\n",
    "from functional import seq\n",
    "import re, os\n",
    "from collections import namedtuple\n",
    "from lxml import etree\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### LeetCode"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "github爬取的md题目有两类，一类是以#开头的，一类是以`<h2>`开头的\n",
    "\n",
    "h2处理报错之后调用该函数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "LeetCodeQuestionData = namedtuple('LeetCodeQuestion', ['title', 'question_text'])\n",
    "\n",
    "def extract_lines(content):\n",
    "    tree = etree.fromstring(content, etree.HTMLParser())\n",
    "    lines = tree.xpath(\"//text()\")\n",
    "    return lines\n",
    "\n",
    "def extract_question(content):\n",
    "    question_lines = []\n",
    "    for line in extract_lines(content):\n",
    "        if \"example\" in line.lower():\n",
    "            break\n",
    "        question_lines.append(line)\n",
    "    \n",
    "    # replace consecutive spaces\n",
    "    question_text = re.sub(\"\\s+\", \" \", ' '.join(question_lines))\n",
    "    # remove space before punctuation\n",
    "    question_text = re.sub(\"\\s?([.,!?-])\", r'\\1', question_text)\n",
    "    return question_text\n",
    "\n",
    "def extract_from_html(content):\n",
    "    title_pattern = re.compile(\"(<h2>.*?</h2><h3>.*?</h3>)\", re.S)\n",
    "    title = title_pattern.findall(content)[0]\n",
    "    title, level = extract_lines(title)\n",
    "\n",
    "    content = \"<body>{}</body>\".format(content)\n",
    "    content = re.sub(title_pattern, '', content)\n",
    "    \n",
    "    question_text = extract_question(content)\n",
    "    return LeetCodeQuestionData(title, question_text)\n",
    "\n",
    "def extract_from_html2(content):\n",
    "    title_pattern = re.compile(\"(<h2>.*?</h2>)\", re.S)\n",
    "    title = title_pattern.findall(content)[0]\n",
    "    title = extract_lines(title)[0]\n",
    "\n",
    "    content = \"<body>{}</body>\".format(content)\n",
    "    content = re.sub(title_pattern, '', content)\n",
    "    \n",
    "    question_text = extract_question(content)\n",
    "    return LeetCodeQuestionData(title, question_text)\n",
    "\n",
    "def extract_from_md(content):\n",
    "    title = re.findall(\"#\\s(.*?)\\n\", content)[0]\n",
    "    content = re.sub(\"#.*\", \"\", content)\n",
    "    content = \"<body>{}</body>\".format(content)\n",
    "    \n",
    "    question_text = extract_question(content)\n",
    "    return LeetCodeQuestionData(title, question_text)\n",
    "\n",
    "def extract(content):\n",
    "    for extractor in [extract_from_html, extract_from_html2, extract_from_md]:\n",
    "        try:\n",
    "            return extractor(content)\n",
    "        except Exception as e:\n",
    "            pass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### leetcode"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# qpaths = glob(\"/home/tanchenwei/Desktop/sematic_search/leetcode/**/README.md\")\n",
    "# qpaths = glob(\"/home/tanchenwei/Desktop/sematic_search/LeetCode/**/README.md\")\n",
    "# qpaths = glob(\"/home/tanchenwei/Desktop/sematic_search/Leetcode-Problems/**/README.md\")\n",
    "\n",
    "def get_all_lcq(qpaths):\n",
    "    qcontents = (seq(qpaths)\n",
    "                .map(lambda x: (x, Path(x).read_text()))\n",
    "                .filter_not(lambda x: 'SQL Schema' in x[1])\n",
    "                .to_list())\n",
    "\n",
    "    lcqs = (seq(qcontents)\n",
    "            .map(lambda x: extract(x[1]))\n",
    "            .filter(lambda x: x is not None)\n",
    "            .map(lambda x: LeetCodeQuestionData(x.title.strip(), x.question_text.strip()))\n",
    "            .to_list())\n",
    "    return lcqs\n",
    "\n",
    "lcqs1 = get_all_lcq(glob(\"/home/tanchenwei/Desktop/sematic_search/LeetCode/**/README.md\"))\n",
    "lcqs2 = get_all_lcq(glob(\"/home/tanchenwei/Desktop/sematic_search/leetcode/**/README.md\"))\n",
    "lcqs3 = get_all_lcq(glob(\"/home/tanchenwei/Desktop/sematic_search/Leetcode-Problems/**/README.md\"))\n",
    "lcqs = lcqs1 + lcqs2 + lcqs3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 94,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(2039, 71)"
      ]
     },
     "execution_count": 94,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "is_ordered_title = lambda x: re.match(\"^\\d+\\.\", x.title)\n",
    "ordered_lcqs = seq(lcqs).filter(is_ordered_title).to_list()\n",
    "unordered_lcqs = seq(lcqs).filter_not(is_ordered_title).to_list()\n",
    "len(ordered_lcqs), len(unordered_lcqs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "metadata": {},
   "outputs": [],
   "source": [
    "ts = (seq(ordered_lcqs)\n",
    "      .map(lambda x: x.title.split(\".\")[1].strip())\n",
    "      .to_list())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 103,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1153 61\n"
     ]
    }
   ],
   "source": [
    "selected = []\n",
    "orders = set()\n",
    "titles = set()\n",
    "\n",
    "for q in ordered_lcqs:\n",
    "    order, title = q.title.split(\".\")\n",
    "    order = int(order)\n",
    "    title = title.strip()\n",
    "    if order in orders:\n",
    "        continue\n",
    "    orders.add(order)\n",
    "    titles.add(title)\n",
    "    selected.append(q)\n",
    "\n",
    "selected_unordered = []\n",
    "for q in unordered_lcqs:\n",
    "    if q.title in titles:\n",
    "        continue\n",
    "    titles.add(q.title)\n",
    "    selected_unordered.append(q)\n",
    "print(len(selected), len(selected_unordered))\n",
    "final_lcq = selected + selected_unordered"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 106,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "with open(\"lcq.pkl\", \"wb\") as f:\n",
    "    pickle.dump(seq(final_lcq)\n",
    "                .map(lambda x: (x.title, x.question_text))\n",
    "                .to_list() , f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 107,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "516K\tlcq.pkl\n"
     ]
    }
   ],
   "source": [
    "!du -h lcq.pkl"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### feature and search"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 108,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/tanchenwei/miniconda3/lib/python3.11/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "# extract feature of question\n",
    "from sentence_transformers import SentenceTransformer, util\n",
    "model_dir = '/home/tanchenwei/all-mpnet-base-v2'\n",
    "model = SentenceTransformer(model_dir)\n",
    "\n",
    "def extract_feature(sentences):\n",
    "    # 获取句子表示\n",
    "    sentence_embeddings = model.encode(sentences)\n",
    "    return sentence_embeddings\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 113,
   "metadata": {},
   "outputs": [],
   "source": [
    "lcq = pickle.load(open(\"lcq.pkl\", \"rb\"))\n",
    "features = extract_feature(seq(lcq)\n",
    "                           .map(lambda x: x[1])\n",
    "                           .to_list())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 114,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1214, 768)"
      ]
     },
     "execution_count": 114,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "features.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 125,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "np.save(\"lcq.npy\", features)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [],
   "source": [
    "target_feature = extract_feature([\"The high number test is finished, and it is difficult for the math teacher to do. The math teacher wants to get out of the students, she always has to increase the score (all positive integers) for some students over and over again, and pay attention to what the lowest score is. Can you help her because of the heavy workload?\"])\n",
    "cos_sim = util.pytorch_cos_sim(target_feature, features)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(75)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cos_sim[0].argmax()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[('1283. Find the Smallest Divisor Given a Threshold', tensor(0.4429)),\n",
       " ('1770. Maximum Score from Performing Multiplication Operations',\n",
       "  tensor(0.4061)),\n",
       " ('202. Happy Number', tensor(0.3992)),\n",
       " ('740. Delete and Earn', tensor(0.3780)),\n",
       " ('167. Two Sum II - Input Array Is Sorted', tensor(0.3769)),\n",
       " ('1291. Sequential Digits', tensor(0.3754)),\n",
       " ('552. Student Attendance Record II', tensor(0.3751)),\n",
       " ('991. Broken Calculator', tensor(0.3674)),\n",
       " ('287. Find the Duplicate Number', tensor(0.3462)),\n",
       " ('287. Find the Duplicate Number', tensor(0.3462))]"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(seq(enumerate(cos_sim[0]))\n",
    " .sorted(key=lambda x: x[1], reverse=True)\n",
    " .map(lambda x: (lcqs[x[0]].title, x[1]))\n",
    " .to_list())[:10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'scp tanchenwei@benke:/home/tanchenwei/Desktop/sematic_search/READMD.md .\\nscp tanchenwei@benke:/home/tanchenwei/Desktop/sematic_search/main.ipynb .\\nscp tanchenwei@benke:/home/tanchenwei/Desktop/sematic_search/lcq.pkl .\\nscp tanchenwei@benke:/home/tanchenwei/Desktop/sematic_search/test.py .\\nscp tanchenwei@benke:/home/tanchenwei/Desktop/sematic_search/output.txt .\\nscp tanchenwei@benke:/home/tanchenwei/Desktop/sematic_search/YouDao.py .\\nscp tanchenwei@benke:/home/tanchenwei/Desktop/sematic_search/lcq.npy .\\nscp tanchenwei@benke:/home/tanchenwei/Desktop/sematic_search/main.py .'"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "templ = 'scp tanchenwei@benke:/home/tanchenwei/Desktop/sematic_search/{} .'\n",
    "cmds = []\n",
    "for each in os.listdir('./'):\n",
    "    if os.path.isfile(each):\n",
    "        cmds.append(templ.format(each))\n",
    "'\\n'.join(cmds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
